import math
import paddle
from paddle import nn
import paddle.nn.functional as F
from paddle import ParamAttr
from paddle.nn.initializer import Normal, Constant, KaimingNormal


def get_bias_attr(k):
    stdv = 1.0 / math.sqrt(k * 1.0)
    initializer = paddle.nn.initializer.Uniform(-stdv, stdv)
    bias_attr = ParamAttr(initializer=initializer)
    return bias_attr


class UpSample(nn.Layer):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv2D(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=3,
            padding=1,
            bias_attr=False)
        self.conv_bn1 = nn.BatchNorm(
            num_channels=out_channels,
            param_attr=ParamAttr(
                initializer=paddle.nn.initializer.Constant(value=1.0)),
            bias_attr=ParamAttr(
                initializer=paddle.nn.initializer.Constant(value=1e-4)),
            act='relu')
        self.conv2 = nn.Conv2DTranspose(
            in_channels=out_channels,
            out_channels=out_channels,
            kernel_size=2,
            stride=2,
            weight_attr=ParamAttr(
                initializer=paddle.nn.initializer.KaimingUniform()),
            bias_attr=get_bias_attr(in_channels // 4))
        self.conv_bn2 = nn.BatchNorm(
            num_channels=out_channels,
            param_attr=ParamAttr(
                initializer=paddle.nn.initializer.Constant(value=1.0)),
            bias_attr=ParamAttr(
                initializer=paddle.nn.initializer.Constant(value=1e-4)),
            act="relu")
        self.conv3 = nn.Conv2DTranspose(
            in_channels=out_channels,
            out_channels=out_channels,
            kernel_size=2,
            stride=2,
            weight_attr=ParamAttr(
                initializer=paddle.nn.initializer.KaimingUniform()),
            bias_attr=get_bias_attr(in_channels // 4 ),
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv_bn1(x)
        x = self.conv2(x)
        x = self.conv_bn2(x)
        x = self.conv3(x)
        return x


class FASTHead(nn.Layer):

    def __init__(self, in_channels, k=50, use_thresh=False, **kwargs):
        super().__init__()
        planes = 64
        self.up = UpSample(in_channels, planes)
        self.geo_conv = self._make_conv(planes, planes, 5)
        self.score_conv = self._make_conv(planes, planes, 1, True)
        scale = paddle.create_parameter([2], dtype=paddle.float32, default_initializer=Constant(1.0))
        self.scale = self.add_parameter('scale', scale)
        self.k = k
        self.thresh_conv = None
        if use_thresh:
            self.thresh_conv = self._make_conv(planes, planes, 1)
        
    def _make_conv(self, in_channels, planes,  out_channels, use_focal_loss=False):
        if use_focal_loss:
            weight = ParamAttr(initializer=Normal(0.0, 0.01))
            bias = ParamAttr(initializer=Constant(-4.5))
        else:
            weight = ParamAttr(initializer=KaimingNormal())
            bias = ParamAttr(initializer=Constant(0.0))

        return nn.Sequential(
            nn.Conv2D(in_channels, planes, 3, 1, 1, bias_attr=False),
            nn.BatchNorm2D(planes),
            nn.ReLU(),
            nn.Conv2D(planes, out_channels, 3, 1, 1, weight_attr=weight, bias_attr=bias)
        )

    def step_function(self, x, y):
        return paddle.reciprocal(1 + paddle.exp(-self.k * (x - y)))


    def forward(self, x):
        x = self.up(x)
        score = self.score_conv(x)
        score = F.sigmoid(score)
        if self.thresh_conv and self.training:
            thresh = F.sigmoid(self.thresh_conv(x))
            score = self.step_function(score, thresh)

        geo = self.geo_conv(x)
        x_coord = F.relu(geo[:, :2] * self.scale[0])
        y_coord = F.relu(geo[:, 2:4] * self.scale[1])
        angle = (F.sigmoid(geo[:, 4:]) - 0.5) * math.pi
        geo = paddle.concat([x_coord, y_coord, angle], axis=1)

        return {'f_score': score, 'f_geo': geo}



class FASTHeadSmall(nn.Layer):
    def __init__(self, in_channels, **kwargs):
        super().__init__()
        planes = 64
        self.conv = self._make_inconv(in_channels, planes)
        self.up = self._make_upsample(in_channels, in_channels)

        self.score_conv = nn.Conv2D(planes, 1, 3, 1, 1)
        self.box_conv = nn.Conv2D(planes, 4, 3, 1, 1)
        self.angle_conv = nn.Conv2D(planes, 1, 3, 1, 1)

        scale = paddle.create_parameter([2], dtype=paddle.float32, 
            default_initializer=nn.initializer.Constant(1.0))
        self.scale = self.add_parameter('scale', scale)

    def _make_upsample(self, in_channels, out_channels):
        return nn.Sequential(
            nn.Conv2DTranspose(in_channels, out_channels, 2, 2, bias_attr=False),
            nn.BatchNorm2D(out_channels),
            nn.ReLU()
        )

    def _make_inconv(self, in_channels, out_channels):
        return nn.Sequential(
            nn.Conv2D(in_channels, in_channels, 3, 1, 1, bias_attr=False),
            nn.BatchNorm2D(in_channels),
            nn.ReLU(),
            nn.Conv2D(in_channels, out_channels, 3, 1, 1)
        )

    def _make_conv(self, in_channels, planes,  out_channels):
        return nn.Sequential(
            nn.Conv2D(in_channels, planes, 3, 1, 1, bias_attr=False),
            nn.BatchNorm2D(planes),
            nn.ReLU(),
            nn.Conv2D(planes, out_channels, 3, 1, 1)
        )
        
    def forward(self, x):
        x = self.up(x)
        x = self.conv(x)
        score = self.score_conv(x)
        box = self.box_conv(x)
        angle = self.angle_conv(x)

        score = F.sigmoid(score)
        x_coord = F.relu(box[:, :2] * self.scale[0])
        y_coord = F.relu(box[:, 2:] * self.scale[1])
        # print(self.scale)

        angle = (F.sigmoid(angle) - 0.5) * math.pi
        geo = paddle.concat([x_coord, y_coord, angle], axis=1)
        return {'f_score': score, 'f_geo': geo}